import numpy as np
import cv2

image = cv2.imread('data/mine.BMP')
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
gray = cv2.medianBlur(gray, 15)

edges = cv2.Canny(gray, 50, 70)

edges = cv2.dilate(edges, None, iterations=1)
edges = cv2.erode(edges, None, iterations=1)

contours, hierarchy = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

position = []
for cnt in contours:
    (x, y, w, h) = cv2.boundingRect(cnt)
    area = cv2.contourArea(cnt)
    perimeter = cv2.arcLength(cnt, True)

    if (area >= 20) & (perimeter > 50) & (w - h <= 10) & (h - w <= 10):
        position.append([x, y, x + w, y + h])


object_3d_points = np.array(([75, -75, 0],
                             [-75, -75, 0],
                             [75, 75, 0],
                             [-75, 75, 0]), dtype=np.double)
object_2d_point = np.array(([position[0][2], position[0][3]],
                            [position[1][0], position[1][3]],
                            [position[2][2], position[2][1]],
                            [position[3][0], position[3][1]]), dtype=np.double)
camera_matrix = np.array(([[1.80212858e+03, 0.00000000e+00, 6.47641612e+02],
                           [0.00000000e+00, 1.80170930e+03, 4.76919503e+02],
                           [0.00000000e+00, 0.00000000e+00, 1.00000000e+00]]), dtype=np.double)
dist_coefs = np.array([-6.95203039e-02, 4.79728273e-01, -6.47191393e-04, -1.01347865e-04, -6.52708458e-01],
                      dtype=np.double)
# 求解相机位姿
found, rvec, tvec = cv2.solvePnP(object_3d_points, object_2d_point, camera_matrix, dist_coefs)
rotM = cv2.Rodrigues(rvec)[0]
camera_postion = -np.matrix(rotM).T * np.matrix(tvec)
print(camera_postion.T)

originpoints, _ = cv2.projectPoints((0, 0, 0), rotM, tvec, camera_matrix, dist_coefs)
zpoints, _ = cv2.projectPoints((0, 0, 100), rotM, tvec, camera_matrix, dist_coefs)
xpoints, _ = cv2.projectPoints((200, 0, 0), rotM, tvec, camera_matrix, dist_coefs)
ypoints, _ = cv2.projectPoints((0, 200, 0), rotM, tvec, camera_matrix, dist_coefs)

imgpoints = np.array(zpoints).astype(np.int16)
originpoints = np.array(originpoints).astype(np.int16)
xpoints = np.array(xpoints).astype(np.int16)
ypoints = np.array(ypoints).astype(np.int16)

cv2.line(image, (originpoints[0][0][0], originpoints[0][0][1]), (imgpoints[0][0][0], imgpoints[0][0][1]), (255, 0, 0), 3, cv2.LINE_8)
cv2.line(image, (originpoints[0][0][0], originpoints[0][0][1]), (xpoints[0][0][0], xpoints[0][0][1]), (0, 255, 0), 3, cv2.LINE_8)
cv2.line(image, (originpoints[0][0][0], originpoints[0][0][1]), (ypoints[0][0][0], ypoints[0][0][1]), (0, 0, 255), 3, cv2.LINE_8)
cv2.circle(image, (originpoints[0][0][0], originpoints[0][0][1]), 4, (255, 255, 255), -1)
str = "dist=" + str(float(camera_postion[2][0])) + "mm"
cv2.putText(image, str, (originpoints[0][0][0], originpoints[0][0][1]), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 255, 0), 2)

cv2.imshow("img", image)
cv2.waitKey()
cv2.destroyAllWindows()

